# src/image_processing/image_analysis.py

import logging
from PIL import Image
import numpy as np
from scipy.stats import entropy, pearsonr, chisquare
import matplotlib.pyplot as plt
import os
import cv2


def plot_histogram_and_chi_square(img_array, mode, title, save_path, bins=256):
    """
    绘制并保存图像的直方图，并计算卡方检验
    :param img_array: 图像数组（numpy 数组）
    :param mode: 图像模式（如 'L', 'RGB', 'RGBA'）
    :param title: 图像标题（包含算法信息）
    :param save_path: 直方图保存路径
    :param bins: 直方图的区间数
    :return: 直方图文件名和卡方检验结果
    """
    plt.figure(figsize=(10, 5))
    chi_square_results = []

    if mode == 'L':  # 灰度图像
        hist, _ = np.histogram(img_array.flatten(), bins=bins, range=(0, 256))
        plt.hist(img_array.flatten(), bins=bins, range=(0, 256), color='black', alpha=0.7)
        plt.title(f'{title} - Gray Channel')

        # 计算卡方检验
        expected = np.full_like(hist, fill_value=np.mean(hist))
        expected = expected * (np.sum(hist) / np.sum(expected))
        try:
            chi_stat, p_value = chisquare(f_obs=hist, f_exp=expected)
        except ValueError as e:
            logging.error(f"卡方检验失败 (Gray Channel): {e}")
            chi_stat, p_value = None, None
        chi_square_results.append((chi_stat, p_value))

    elif mode == 'RGB':  # RGB 图像
        colors = ['red', 'green', 'blue']
        channel_names = ['Red', 'Green', 'Blue']
        for i in range(3):
            hist, _ = np.histogram(img_array[:, :, i].flatten(), bins=bins, range=(0, 256))
            plt.hist(img_array[:, :, i].flatten(), bins=bins, range=(0, 256), color=colors[i], alpha=0.7,
                     label=channel_names[i])

            # 计算卡方检验
            expected = np.full_like(hist, fill_value=np.mean(hist))
            expected = expected * (np.sum(hist) / np.sum(expected))
            try:
                chi_stat, p_value = chisquare(f_obs=hist, f_exp=expected)
            except ValueError as e:
                logging.error(f"卡方检验失败 (Channel {i}): {e}")
                chi_stat, p_value = None, None
            chi_square_results.append((chi_stat, p_value))

        plt.title(f'{title} - RGB Channels')
        plt.legend()

    elif mode == 'RGBA':  # RGBA 图像
        colors = ['red', 'green', 'blue', 'gray']
        channel_names = ['Red', 'Green', 'Blue', 'Alpha']
        for i in range(4):
            hist, _ = np.histogram(img_array[:, :, i].flatten(), bins=bins, range=(0, 256))
            plt.hist(img_array[:, :, i].flatten(), bins=bins, range=(0, 256), color=colors[i], alpha=0.7,
                     label=channel_names[i])

            # 计算卡方检验
            expected = np.full_like(hist, fill_value=np.mean(hist))
            expected = expected * (np.sum(hist) / np.sum(expected))
            try:
                chi_stat, p_value = chisquare(f_obs=hist, f_exp=expected)
            except ValueError as e:
                logging.error(f"卡方检验失败 (Channel {i}): {e}")
                chi_stat, p_value = None, None
            chi_square_results.append((chi_stat, p_value))

        plt.title(f'{title} - RGBA Channels')
        plt.legend()

    plt.xlabel('Pixel Value')
    plt.ylabel('Frequency')

    # 生成直方图文件名
    histogram_filename = f"{title.replace(' ', '_')}_histogram.png"
    histogram_path = os.path.join(save_path, histogram_filename)

    plt.savefig(histogram_path)
    plt.close()

    return histogram_filename, chi_square_results


def calculate_correlation(img_array, mode):
    """
    计算图像的相关性（水平、垂直、对角线）
    :param img_array: 图像数组
    :param mode: 图像模式
    :return: 相关性分析结果
    """
    correlations = []
    if mode == 'L':  # 灰度图像
        # 水平相关性
        x_horizontal = img_array[:, :-1].flatten()
        y_horizontal = img_array[:, 1:].flatten()
        if np.all(x_horizontal == x_horizontal[0]) or np.all(y_horizontal == y_horizontal[0]):
            corr_horizontal = 0
        else:
            corr_horizontal, _ = pearsonr(x_horizontal, y_horizontal)

        # 垂直相关性
        x_vertical = img_array[:-1, :].flatten()
        y_vertical = img_array[1:, :].flatten()
        if np.all(x_vertical == x_vertical[0]) or np.all(y_vertical == y_vertical[0]):
            corr_vertical = 0
        else:
            corr_vertical, _ = pearsonr(x_vertical, y_vertical)

        # 对角相关性
        x_diagonal = img_array[:-1, :-1].flatten()
        y_diagonal = img_array[1:, 1:].flatten()
        if np.all(x_diagonal == x_diagonal[0]) or np.all(y_diagonal == y_diagonal[0]):
            corr_diagonal = 0
        else:
            corr_diagonal, _ = pearsonr(x_diagonal, y_diagonal)

        correlations.append((corr_horizontal, corr_vertical, corr_diagonal))
    else:  # 多通道图像
        for i in range(img_array.shape[2]):
            # 水平相关性
            x_horizontal = img_array[:, :-1, i].flatten()
            y_horizontal = img_array[:, 1:, i].flatten()
            if np.all(x_horizontal == x_horizontal[0]) or np.all(y_horizontal == y_horizontal[0]):
                corr_horizontal = 0
            else:
                corr_horizontal, _ = pearsonr(x_horizontal, y_horizontal)

            # 垂直相关性
            x_vertical = img_array[:-1, :, i].flatten()
            y_vertical = img_array[1:, :, i].flatten()
            if np.all(x_vertical == x_vertical[0]) or np.all(y_vertical == y_vertical[0]):
                corr_vertical = 0
            else:
                corr_vertical, _ = pearsonr(x_vertical, y_vertical)

            # 对角相关性
            x_diagonal = img_array[:-1, :-1, i].flatten()
            y_diagonal = img_array[1:, 1:, i].flatten()
            if np.all(x_diagonal == x_diagonal[0]) or np.all(y_diagonal == y_diagonal[0]):
                corr_diagonal = 0
            else:
                corr_diagonal, _ = pearsonr(x_diagonal, y_diagonal)

            correlations.append((corr_horizontal, corr_vertical, corr_diagonal))
    return correlations


def calculate_entropy(img_array, mode):
    """
    计算图像的信息熵
    :param img_array: 图像数组
    :param mode: 图像模式
    :return: 各通道的信息熵
    """
    entropy_values = []
    if mode == 'L':  # 灰度图像
        hist, _ = np.histogram(img_array.flatten(), bins=256, range=(0, 256))
        prob = hist / np.sum(hist)
        entropy_values.append(entropy(prob, base=2))
    else:  # 多通道图像
        for i in range(img_array.shape[2]):
            hist, _ = np.histogram(img_array[:, :, i].flatten(), bins=256, range=(0, 256))
            prob = hist / np.sum(hist)
            entropy_values.append(entropy(prob, base=2))
    return entropy_values


def calculate_statistics(img_array, mode):
    """
    计算图像的统计信息（均值、方差）
    :param img_array: 图像数组
    :param mode: 图像模式
    :return: 各通道的均值和方差
    """
    statistics = []
    if mode == 'L':  # 灰度图像
        mean = np.mean(img_array)
        var = np.var(img_array)
        statistics.append((mean, var))
    else:  # 多通道图像
        for i in range(img_array.shape[2]):
            mean = np.mean(img_array[:, :, i])
            var = np.var(img_array[:, :, i])
            statistics.append((mean, var))
    return statistics


def calculate_psnr(original_img_array, encrypted_img_array, mode):
    """
    计算峰值信噪比（PSNR）
    :param original_img_array: 原始图像数组
    :param encrypted_img_array: 加密图像数组
    :param mode: 图像模式
    :return: PSNR值
    """
    if mode == 'L':  # 灰度图像
        mse = np.mean((original_img_array - encrypted_img_array) ** 2)
        if mse == 0:
            return [float('inf')]
        psnr = 10 * np.log10((255 ** 2) / mse)
        return [psnr]
    else:  # 多通道图像
        psnr_values = []
        for i in range(original_img_array.shape[2]):
            mse = np.mean((original_img_array[:, :, i] - encrypted_img_array[:, :, i]) ** 2)
            if mse == 0:
                psnr_values.append(float('inf'))
            else:
                psnr = 20 * np.log10(255 / np.sqrt(mse))
                psnr_values.append(psnr)
        return psnr_values


def calculate_gvd(original_img_array, encrypted_img_array, mode):
    """
    计算梯度变化度（GVD）
    :param original_img_array: 原始图像数组
    :param encrypted_img_array: 加密图像数组
    :param mode: 图像模式
    :return: GVD 值（单通道或多通道）
    """
    if mode == 'L':  # 灰度图像
        original_gray = original_img_array.astype(np.int16)
        encrypted_gray = encrypted_img_array.astype(np.int16)
        gvd = GVD_components(original_gray, encrypted_gray)
        return [gvd]
    elif mode == 'RGB':  # RGB 图像
        B1, G1, R1 = cv2.split(original_img_array)
        B2, G2, R2 = cv2.split(encrypted_img_array)
        R1 = R1.astype(np.int16)
        R2 = R2.astype(np.int16)
        G1 = G1.astype(np.int16)
        G2 = G2.astype(np.int16)
        B1 = B1.astype(np.int16)
        B2 = B2.astype(np.int16)
        R_gvd = GVD_components(R1, R2)
        G_gvd = GVD_components(G1, G2)
        B_gvd = GVD_components(B1, B2)
        return [R_gvd, G_gvd, B_gvd]
    elif mode == 'RGBA':  # RGBA 图像
        B1, G1, R1, A1 = cv2.split(original_img_array)
        B2, G2, R2, A2 = cv2.split(encrypted_img_array)
        R1 = R1.astype(np.int16)
        R2 = R2.astype(np.int16)
        G1 = G1.astype(np.int16)
        G2 = G2.astype(np.int16)
        B1 = B1.astype(np.int16)
        B2 = B2.astype(np.int16)
        A1 = A1.astype(np.int16)
        A2 = A2.astype(np.int16)
        R_gvd = GVD_components(R1, R2)
        G_gvd = GVD_components(G1, G2)
        B_gvd = GVD_components(B1, B2)
        A_gvd = GVD_components(A1, A2)
        return [R_gvd, G_gvd, B_gvd, A_gvd]
    else:
        raise ValueError(f"不支持的图像模式: {mode}")


def GVD_components(channel0, channel1):
    """
    计算图像分量的 GVD
    :param channel0: 原始图像分量（单通道）
    :param channel1: 加密图像分量（单通道）
    :return: GVD 值
    """
    if len(channel0.shape) != 2 or len(channel1.shape) != 2:
        raise ValueError("GVD_components 函数只支持单通道图像输入")

    AN0 = 0  # 原始图像的梯度总和
    AN1 = 0  # 加密图像的梯度总和
    w, h = channel0.shape  # 获取图像的宽度和高度
    for i in range(1, w - 1):  # 遍历图像的每个像素（排除边缘）
        for j in range(1, h - 1):  # 遍历图像的每个像素（排除边缘）
            # 加密前
            GN0 = 0  # 原始图像的梯度
            GN0 += (channel0[i][j] - channel0[i - 1][j]) ** 2  # 计算垂直梯度
            GN0 += (channel0[i][j] - channel0[i + 1][j]) ** 2  # 计算垂直梯度
            GN0 += (channel0[i][j] - channel0[i][j - 1]) ** 2  # 计算水平梯度
            GN0 += (channel0[i][j] - channel0[i][j + 1]) ** 2  # 计算水平梯度
            GN0 /= 4  # 计算平均梯度

            # 加密后
            GN1 = 0  # 加密图像的梯度
            GN1 += (channel1[i][j] - channel1[i - 1][j]) ** 2  # 计算垂直梯度
            GN1 += (channel1[i][j] - channel1[i + 1][j]) ** 2  # 计算垂直梯度
            GN1 += (channel1[i][j] - channel1[i][j - 1]) ** 2  # 计算水平梯度
            GN1 += (channel1[i][j] - channel1[i][j + 1]) ** 2  # 计算水平梯度
            GN1 /= 4  # 计算平均梯度

            AN0 += GN0  # 累加原始图像的梯度
            AN1 += GN1  # 累加加密图像的梯度

    AN0 /= ((w - 2) * (h - 2))  # 计算原始图像的平均梯度
    AN1 /= ((w - 2) * (h - 2))  # 计算加密图像的平均梯度

    # 处理分母为零的情况
    if AN1 + AN0 == 0:
        return 0  # 如果分母为零，返回 0
    gvd = (AN1 - AN0) / (AN1 + AN0)  # 计算 GVD
    return gvd  # 返回 GVD 值


def analyze_image(image_data, title, save_dir, original_image_data=None):
    """
    分析图像
    :param image_data: ImageData 对象
    :param title: 图像标题（包含算法信息）
    :param save_dir: 保存目录
    :param original_image_data: 原始图像的 ImageData 对象（用于计算PSNR和GVD）
    :return: 分析结果
    """
    img_array = image_data.get_image_array()  # 获取图像数组
    mode = image_data.get_mode()  # 获取图像模式

    # 确保保存目录存在
    os.makedirs(save_dir, exist_ok=True)

    #TODO: 效率优化——原始图像分析也进行了卡方检验，这是应该避免的！

    # 绘制并保存直方图，并计算卡方检验
    histogram_filename, chi_square_results = plot_histogram_and_chi_square(img_array, mode, title, save_dir)

    # 计算相关性、信息熵、统计信息
    correlations = calculate_correlation(img_array, mode)  # 计算相关性
    entropy_values = calculate_entropy(img_array, mode)  # 计算信息熵
    statistics = calculate_statistics(img_array, mode)  # 计算统计信息

    # 计算 PSNR
    if original_image_data is not None:  # 如果有原始图像数据
        original_img_array = original_image_data.get_image_array()  # 获取原始图像数组
        psnr_values = calculate_psnr(original_img_array, img_array, mode)  # 计算 PSNR
    else:  # 如果没有原始图像数据
        psnr_values = [None] * len(entropy_values)  # 返回 None

    # 返回分析结果
    return {
        'title': title,  # 图像标题
        'mode': mode,  # 图像模式
        'histogram_path': histogram_filename,  # 直方图路径
        'correlations': correlations,  # 相关性结果
        'entropy': entropy_values,  # 信息熵结果
        'statistics': statistics,  # 统计信息结果
        'chi_square': chi_square_results,  # 卡方检验结果
        'psnr': psnr_values  # PSNR 结果
    }
